import numpy as np
import scipy as sp
from osqp_demo import OSQP
from scipy import sparse


class MPC:
    n: int
    nx: int
    nu: int
    Q: sparse.csc_matrix
    R: sparse.csc_matrix
    A: list[sparse.csc_matrix]
    B: list[sparse.csc_matrix]
    xMin: np.ndarray
    xMax: np.ndarray
    uMin: np.ndarray
    uMax: np.ndarray

    def __init__(self, n: int, nx: int, nu: int) -> None:
        self.n = n
        self.nx = nx
        self.nu = nu

    def setQ(self, Q: np.ndarray) -> None:
        self.Q = sparse.csc_matrix(Q)

    def setR(self, R: np.ndarray) -> None:
        self.R = sparse.csc_matrix(R)

    def setA(self, As: list[np.ndarray]) -> None:
        self.A = [sparse.csc_matrix(A) for A in As]

    def setB(self, Bs: list[np.ndarray]) -> None:
        self.B = [sparse.csc_matrix(B) for B in Bs]

    def setStateRange(self, min: np.ndarray, max: np.ndarray) -> None:
        self.xMin = min
        self.xMax = max

    def setInputRange(self, min: np.ndarray, max: np.ndarray) -> None:
        self.uMin = min
        self.uMax = max

    # xRef shape((n+1)*nx) x0,x1,x2,x3--xn
    # xCur shape(nx) x0
    def solve(
        self, xRef: np.ndarray, xCur: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray]:
        P = sparse.block_diag(
            [
                sparse.kron(sparse.eye(self.n + 1), self.Q),
                sparse.kron(sparse.eye(self.n), self.R),
            ]
        )
        c = np.kron(np.ones((self.n + 1, 1)), -self.Q.toarray())
        # qT  但OSQP库实际上是否转置都可以
        q = np.vstack(
            [
                np.kron(np.eye(self.n + 1), -self.Q.toarray()) @ xRef,
                np.zeros((self.n * self.nu, 1)),
            ]
        )
        ADynamic = sparse.hstack(
            [
                sparse.hstack(
                    [
                        sparse.vstack(
                            [
                                sparse.csc_matrix((self.nx, self.n * self.nx)),
                                sparse.block_diag(self.A),
                            ]
                        ),
                        sparse.csc_matrix(((self.n + 1) * self.nx, self.nx)),
                    ]
                )
                - sparse.kron(sparse.eye(self.n + 1), sparse.eye(self.nx)),
                sparse.vstack(
                    [
                        sparse.csc_matrix((self.nx, self.n * self.nu)),
                        sparse.block_diag(self.B),
                    ]
                ),
            ]
        )
        lowerDynamic = np.vstack([-xCur, np.zeros((self.n * self.nx, 1))])
        upperDynamic = np.vstack([-xCur, np.zeros((self.n * self.nx, 1))])
        A = sparse.vstack(
            [ADynamic, sparse.eye((self.n + 1) * self.nx + self.n * self.nu)]
        )
        l = np.vstack(
            [
                lowerDynamic,
                np.kron(np.ones((self.n + 1, 1)), self.xMin),
                np.kron(np.ones((self.n, 1)), self.uMin),
            ]
        )
        u = np.vstack(
            [
                upperDynamic,
                np.kron(np.ones((self.n + 1, 1)), self.xMax),
                np.kron(np.ones((self.n, 1)), self.uMax),
            ]
        )
        osqp = OSQP()
        osqp.setup(P=P, q=q, A=A, l=l, u=u)
        res = osqp.solve()
        if res.info.status != "solved":
            return np.array([]), np.array([])
        states: np.ndarray = res.x[: (self.n + 1) * self.nx].reshape(
            self.n + 1, self.nx
        )
        controls: np.ndarray = res.x[(self.n + 1) * self.nx :].reshape(self.n, self.nu)
        return states, controls
